package com.example.scribe.backend.handlers;

import com.example.scribe.SyntaxNode;
import com.example.scribe.backend.ASTUtils;
import com.example.scribe.backend.BytecodeOps;
import com.example.scribe.backend.ExpressionEmitter;
import com.example.scribe.backend.LocalVariableAllocator;
import com.example.scribe.backend.StatementEmitter;
import com.example.scribe.backend.StatementNodeHandler;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;

/**
 * Handler for assignment statements.
 * Generates bytecode for: identifier = expression
 * Supports: =, +=, -=, *=, /=, %=, <<=, >>=, >>>=, &=, |=, ^=
 */
public final class AssignmentStatementHandler implements StatementNodeHandler {
    public AssignmentStatementHandler() {
    }

    @Override
    public boolean canHandle(String nodeType) {
        return "assignment".equals(nodeType);
    }

    @Override
    public StackManipulation handle(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals, ExpressionEmitter exprEmitter, StatementEmitter statementEmitter) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }
        
        // Assignment might be wrapped in expression_stmt, so check for assignment node
        SyntaxNode assignmentNode = node;
        if (!"assignment".equals(node.type())) {
            // Try to find assignment node in children
            assignmentNode = ASTUtils.findFirst(node, "assignment");
            if (assignmentNode == null) {
                return StackManipulation.Trivial.INSTANCE;
            }
        }
        
        if (assignmentNode.children().size() < 2) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Assignment structure: conditional '=' assignment
        // For simple assignment like "i = i + 1", we need to:
        // 1. Find the left-hand side (identifier)
        // 2. Find the right-hand side (expression after '=')
        // 3. Emit right-hand side expression
        // 4. Store result to left-hand side variable

        SyntaxNode leftNode = null;
        SyntaxNode rightNode = null;
        String op = "=";

        // Assignment node structure: conditional ('=' | '+=' | '-=') assignment
        // For "i = i + 1", the structure is:
        // assignment
        //   conditional (contains identifier "i")
        //   Terminal: '='
        //   assignment (contains "i + 1")
        //     conditional (contains "i")
        //     addition (contains "+")
        //     conditional (contains "1")
        
        // Find the '=' operator (Terminal node)
        int equalsIndex = -1;
        for (int i = 0; i < assignmentNode.children().size(); i++) {
            SyntaxNode child = assignmentNode.children().get(i);
            String text = child.text();
            String type = child.type();
            // Check both direct text match and Terminal node
            // Support: =, +=, -=, *=, /=, %=, <<=, >>=, >>>=, &=, |=, ^=
            if (("=".equals(text) || "+=".equals(text) || "-=".equals(text) || "*=".equals(text) || "/=".equals(text) ||
                 "%=".equals(text) || "<<=".equals(text) || ">>=".equals(text) || ">>>=".equals(text) ||
                 "&=".equals(text) || "|=".equals(text) || "^=".equals(text)) || 
                (type.startsWith("Terminal:") && ("=".equals(text) || "+=".equals(text) || "-=".equals(text) || "*=".equals(text) || "/=".equals(text) ||
                 "%=".equals(text) || "<<=".equals(text) || ">>=".equals(text) || ">>>=".equals(text) ||
                 "&=".equals(text) || "|=".equals(text) || "^=".equals(text)))) {
                op = text;
                equalsIndex = i;
                // Left side is before operator (conditional)
                if (i > 0) {
                    leftNode = assignmentNode.children().get(i - 1);
                }
                // Right side is after operator (assignment or conditional)
                if (i + 1 < assignmentNode.children().size()) {
                    rightNode = assignmentNode.children().get(i + 1);
                }
                break;
            }
        }

        // If not found in direct children, search recursively for Terminal with '='
        if (equalsIndex == -1) {
            for (int i = 0; i < assignmentNode.children().size(); i++) {
                SyntaxNode child = assignmentNode.children().get(i);
                if (child.type().startsWith("Terminal:")) {
                    String text = child.text();
                    if ("=".equals(text) || "+=".equals(text) || "-=".equals(text) || "*=".equals(text) || "/=".equals(text) ||
                        "%=".equals(text) || "<<=".equals(text) || ">>=".equals(text) || ">>>=".equals(text) ||
                        "&=".equals(text) || "|=".equals(text) || "^=".equals(text)) {
                        op = text;
                        equalsIndex = i;
                        if (i > 0) {
                            leftNode = assignmentNode.children().get(i - 1);
                        }
                        if (i + 1 < assignmentNode.children().size()) {
                            rightNode = assignmentNode.children().get(i + 1);
                        }
                        break;
                    }
                }
            }
        }

        // If still not found, try alternative structure:
        // assignment might have: [conditional, assignment] where first is left, second is right
        if (leftNode == null || rightNode == null) {
            // First child should be conditional (contains identifier)
            if (assignmentNode.children().size() >= 2) {
                SyntaxNode first = assignmentNode.children().get(0);
                // Search for identifier in first child
                SyntaxNode found = ASTUtils.findFirst(first, "identifier");
                if (found != null && ASTUtils.isIdentifier(found)) {
                    leftNode = found;
                } else if (ASTUtils.isIdentifier(first)) {
                    leftNode = first;
                } else if ("conditional".equals(first.type())) {
                    // Use the whole conditional as left (will extract identifier from it)
                    leftNode = first;
                }
                
                // Last child should be assignment or conditional (contains right-hand expression)
                SyntaxNode last = assignmentNode.children().get(assignmentNode.children().size() - 1);
                if ("assignment".equals(last.type()) || "conditional".equals(last.type()) || "expression".equals(last.type())) {
                    rightNode = last;
                } else {
                    // Try to find expression in children
                    SyntaxNode expr = ASTUtils.findFirst(last, "expression");
                    if (expr != null) {
                        rightNode = expr;
                    }
                }
            }
        }

        if (leftNode == null || rightNode == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Extract identifier name from left node
        // leftNode might be conditional, primary, or direct identifier
        String varName = null;
        if (ASTUtils.isIdentifier(leftNode)) {
            varName = leftNode.text();
        } else {
            // Search recursively for identifier
            SyntaxNode idNode = ASTUtils.findFirst(leftNode, "identifier");
            if (idNode != null && ASTUtils.isIdentifier(idNode)) {
                varName = idNode.text();
            } else {
                // Try direct children
                for (SyntaxNode child : leftNode.children()) {
                    if (ASTUtils.isIdentifier(child)) {
                        varName = child.text();
                        break;
                    }
                    // Also check in primary nodes
                    if ("primary".equals(child.type())) {
                        for (SyntaxNode prim : child.children()) {
                            if (ASTUtils.isIdentifier(prim)) {
                                varName = prim.text();
                                break;
                            }
                        }
                        if (varName != null) break;
                    }
                }
            }
        }

        if (varName == null || !locals.hasName(varName)) {
            return StackManipulation.Trivial.INSTANCE;
        }

        int varIndex = locals.getIndexByName(varName);

        // Handle compound assignment operators: +=, -=, *=, /=, %=, <<=, >>=, >>>=, &=, |=, ^=
        // For compound assignment, we need to: load variable, emit right expression, perform operation, store
        if ("+=".equals(op) || "-=".equals(op) || "*=".equals(op) || "/=".equals(op) ||
            "%=".equals(op) || "<<=".equals(op) || ">>=".equals(op) || ">>>=".equals(op) ||
            "&=".equals(op) || "|=".equals(op) || "^=".equals(op)) {
            // Load the variable value
            StackManipulation loadVar = BytecodeOps.loadInt(varIndex);
            
            // Emit right-hand side expression
            StackManipulation rightEmit = exprEmitter.emit(rightNode, generatedClassName, locals);
            if (rightEmit == StackManipulation.Trivial.INSTANCE) {
                // Fallback: try to find identifier or number in right node
                if (ASTUtils.isNumber(rightNode)) {
                    try {
                        int v = Integer.parseInt(rightNode.text());
                        rightEmit = BytecodeOps.pushInt(v);
                    } catch (NumberFormatException e) {
                        return StackManipulation.Trivial.INSTANCE;
                    }
                } else if (ASTUtils.isIdentifier(rightNode)) {
                    String rightName = rightNode.text();
                    if (locals.hasName(rightName)) {
                        rightEmit = BytecodeOps.loadInt(locals.getIndexByName(rightName));
                    } else {
                        return StackManipulation.Trivial.INSTANCE;
                    }
                } else {
                    // Try to find expression node in rightNode and emit it
                    SyntaxNode expr = ASTUtils.findFirst(rightNode, "expression");
                    if (expr != null) {
                        rightEmit = exprEmitter.emit(expr, generatedClassName, locals);
                    }
                    if (rightEmit == StackManipulation.Trivial.INSTANCE) {
                        return StackManipulation.Trivial.INSTANCE;
                    }
                }
            }
            
            // Perform the operation based on operator
            StackManipulation operation;
            switch (op) {
                case "+=":
                    operation = BytecodeOps.iadd();
                    break;
                case "-=":
                    operation = BytecodeOps.isub();
                    break;
                case "*=":
                    operation = BytecodeOps.imul();
                    break;
                case "/=":
                    operation = BytecodeOps.idiv();
                    break;
                case "%=":
                    operation = BytecodeOps.irem();
                    break;
                case "<<=":
                    operation = BytecodeOps.ishl();
                    break;
                case ">>=":
                    operation = BytecodeOps.ishr();
                    break;
                case ">>>=":
                    operation = BytecodeOps.iushr();
                    break;
                case "&=":
                    operation = BytecodeOps.iand();
                    break;
                case "|=":
                    operation = BytecodeOps.ior();
                    break;
                case "^=":
                    operation = BytecodeOps.ixor();
                    break;
                default:
                    return StackManipulation.Trivial.INSTANCE;
            }
            
            // Load var, emit right, perform operation, store result
            return new Compound(loadVar, rightEmit, operation, BytecodeOps.storeInt(varIndex));
        }
        
        // Simple assignment (=): just emit right expression and store
        StackManipulation rightEmit = exprEmitter.emit(rightNode, generatedClassName, locals);
        if (rightEmit == StackManipulation.Trivial.INSTANCE) {
            // Fallback: try to find identifier or number in right node
            if (ASTUtils.isNumber(rightNode)) {
                try {
                    int v = Integer.parseInt(rightNode.text());
                    rightEmit = BytecodeOps.pushInt(v);
                } catch (NumberFormatException e) {
                    return StackManipulation.Trivial.INSTANCE;
                }
            } else if (ASTUtils.isIdentifier(rightNode)) {
                String rightName = rightNode.text();
                if (locals.hasName(rightName)) {
                    rightEmit = BytecodeOps.loadInt(locals.getIndexByName(rightName));
                } else {
                    return StackManipulation.Trivial.INSTANCE;
                }
            } else {
                // Try to find expression node in rightNode and emit it
                SyntaxNode expr = ASTUtils.findFirst(rightNode, "expression");
                if (expr != null) {
                    rightEmit = exprEmitter.emit(expr, generatedClassName, locals);
                }
                // If still trivial, try to emit the rightNode itself as expression
                if (rightEmit == StackManipulation.Trivial.INSTANCE && 
                    ("expression".equals(rightNode.type()) || "assignment".equals(rightNode.type()) || "conditional".equals(rightNode.type()))) {
                    // Try one more time with the node itself
                    rightEmit = exprEmitter.emit(rightNode, generatedClassName, locals);
                }
                if (rightEmit == StackManipulation.Trivial.INSTANCE) {
                    return StackManipulation.Trivial.INSTANCE;
                }
            }
        }

        // Store result to variable
        return new Compound(rightEmit, BytecodeOps.storeInt(varIndex));
    }
}

